from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException, ElementClickInterceptedException, ElementNotInteractableException
from selenium.webdriver.remote.webelement import WebElement
import time
import sys
from enum import Enum

# 定义条件类型的枚举
class ConditionType(Enum):
    WIDTH_GREATER_THAN_50 = "width_greater_than_50"
    WIDTH_LESS_THAN_50 = "width_less_than_50"
    ELEMENT_VISIBLE = "element_visible"  # 示例：可扩展更多条件类型

def TemplateFunction(
    driver: webdriver.WebDriver,
    target: str,
    step_name: str = "步骤",
    preconditions: str = "//",
    precondition_type: ConditionType = None,
    postcondition: str = "//",
    postcondition_type: ConditionType = None,
    timeout: int = 10,
    retry_count: int = 0,
    terminate_on_failure: bool = False
) -> WebElement:
    """执行带前置条件和后置条件验证的元素操作，支持自定义条件类型验证。

    Args:
        driver (WebDriver): Selenium WebDriver实例
        target (str): 目标元素的XPath
        step_name (str, optional): 步骤名称
        preconditions (str, optional): 前置条件元素的XPath，"//"表示跳过
        precondition_type (ConditionType, optional): 前置条件类型
        postcondition (str, optional): 后置条件元素的XPath，"//"表示跳过
        postcondition_type (ConditionType, optional): 后置条件类型
        timeout (int, optional): 超时时间（秒）
        retry_count (int, optional): 重试次数
        terminate_on_failure (bool, optional): 失败时是否终止程序

    Returns:
        WebElement: 操作成功返回目标元素，失败返回None
    """
    attempt = 0
    max_attempts = retry_count + 1
    
    while attempt < max_attempts:
        attempt += 1
        print(f"\n提示 : 当前执行 [{step_name}] 步骤，尝试次数 {attempt}/{max_attempts}")
        
        wait = WebDriverWait(driver, timeout)
        element = None
        
        try:
            # 前置条件验证
            if preconditions != "//":
                print(f"提示 : 正在验证 [{step_name}] 前置条件...")
                pre_element = wait.until(
                    EC.presence_of_element_located((By.XPATH, preconditions)),
                    f"前置条件元素未找到: {preconditions}"
                )
                
                if precondition_type:
                    if not _validate_condition(driver, pre_element, precondition_type):
                        raise TimeoutException(f"前置条件验证失败: {precondition_type.value}")
                
                print(f"提示 : [{step_name}] 前置条件验证通过")
            
            # 查找并操作目标元素
            element = wait.until(
                EC.presence_of_element_located((By.XPATH, target)),
                f"目标元素未找到: {target}"
            )
            
            # 滚动到元素位置
            driver.execute_script("arguments[0].scrollIntoView({block: 'center', inline: 'center'});", element)
            time.sleep(0.5)  # 等待滚动完成
            
            # 执行点击操作
            try:
                element.click()
                print(f"提示 : 成功点击目标元素: {target}")
            except (ElementClickInterceptedException, ElementNotInteractableException):
                driver.execute_script("arguments[0].click();", element)
                print(f"提示 : 通过JS点击目标元素: {target}")
            
            # 后置条件验证
            if postcondition != "//":
                print(f"提示 : 正在验证 [{step_name}] 后置条件...")
                post_element = wait.until(
                    EC.presence_of_element_located((By.XPATH, postcondition)),
                    f"后置条件元素未找到: {postcondition}"
                )
                
                if postcondition_type:
                    if not _validate_condition(driver, post_element, postcondition_type):
                        raise TimeoutException(f"后置条件验证失败: {postcondition_type.value}")
                
                print(f"提示 : [{step_name}] 后置条件验证通过，操作成功")
            
            print(f"提示 : [{step_name}] 步骤执行成功！")
            return element
            
        except TimeoutException as te:
            error_msg = f"错误 : [{step_name}] 超时异常: {str(te)}"
            if attempt < max_attempts:
                print(f"警告 : {error_msg}，准备重试...")
                time.sleep(1)  # 重试前等待
            else:
                print(error_msg)
                if terminate_on_failure:
                    sys.exit(1)
                return None
                
        except Exception as e:
            error_msg = f"错误 : [{step_name}] 执行异常: {str(e)}"
            if attempt < max_attempts:
                print(f"警告 : {error_msg}，准备重试...")
                time.sleep(1)  # 重试前等待
            else:
                print(error_msg)
                if terminate_on_failure:
                    sys.exit(1)
                return None

def _validate_condition(driver: webdriver.WebDriver, element: WebElement, condition_type: ConditionType) -> bool:
    """验证元素是否满足指定条件类型"""
    if condition_type == ConditionType.WIDTH_GREATER_THAN_50:
        width = element.size.get('width', 0)
        return width > 50
    elif condition_type == ConditionType.WIDTH_LESS_THAN_50:
        width = element.size.get('width', 0)
        return width < 50
    elif condition_type == ConditionType.ELEMENT_VISIBLE:  # 示例扩展
        return element.is_displayed()
    else:
        print(f"警告 : 未知条件类型: {condition_type}")
        return False    